# 导入必备工具包
import torch  # PyTorch 深度学习框架
import numpy as np  # NumPy，用于处理数组和矩阵操作
from template import *  # 导入自定义模板相关模块
from rich import print  # 导入 rich 库，用于美化打印输出
from datasets import load_dataset  # 导入 Hugging Face 的 datasets 模块，用于加载数据集
from functools import partial  # 导入 partial 函数，用于创建偏函数

from transformers import AutoTokenizer  # 导入 Hugging Face 的 AutoTokenizer，用于文本 tokenization
from data_handle.template import HardTemplate  # 自定义硬模板类
from pet_config import *  # 导入配置文件中的内容（如 ProjectConfig 类）


def convert_example(
        examples: dict,
        tokenizer,
        max_seq_len: int,
        max_label_len: int,
        hard_template: HardTemplate,
        train_mode=True,
        return_tensor=False) -> dict:
    """
    将样本数据转换为模型接收的输入数据。

    Args:
        examples (dict): 训练数据样本, e.g. -> {"text": ['手机\t这个手机也太卡了。', '体育\t世界杯为何迟迟不见宣传']}
        max_seq_len (int): 句子的最大长度，若没有达到最大长度，则 padding 为最大长度
        max_label_len (int): 最大 label 长度，若没有达到最大长度，则 padding 为最大长度
        hard_template (HardTemplate): 模板类。
        train_mode (bool): 训练阶段 or 推理阶段。
        return_tensor (bool): 是否返回 tensor 类型，如不是，则返回 numpy 类型。

    Returns:
        dict (str: np.array) -> tokenized_output = {
                            'input_ids': [[1, 47, 10, 7, 304, 3, 3, 3, 3, 47, 27, 247, 98, 105, 512, 777, 15, 12043, 2], ...],
                            'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], ...],
                            'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], ...],
                            'mask_positions': [[5, 6, 7, 8], ...],
                            'mask_labels': [[2372, 3442, 0, 0], [2643, 4434, 2334, 0], ...]
                        }
    """
    # 初始化存储 tokenized 输出的字典
    tokenized_output = {
        'input_ids': [],
        'token_type_ids': [],
        'attention_mask': [],
        'mask_positions': [],
        'mask_labels': []
    }

    for i, example in enumerate(examples['text']):  # 遍历每个样本
        if train_mode:  # 如果是训练模式
            label, content = example.strip().split('\t')  # 分离标签和内容（以 '\t' 分隔）
        else:  # 如果是推理模式
            content = example.strip()  # 只保留内容部分

        inputs_dict = {  # 构造输入字典
            'textA': content,  # 内容部分
            'MASK': '[MASK]'  # [MASK] 占位符
        }

        encoded_inputs = hard_template(  # 调用硬模板类进行编码
            inputs_dict=inputs_dict,  # 输入字典
            tokenizer=tokenizer,  # Tokenizer 对象
            max_seq_len=max_seq_len,  # 最大序列长度
            mask_length=max_label_len  # MASK 长度
        )

        # 将编码结果添加到 tokenized_output 字典中
        tokenized_output['input_ids'].append(encoded_inputs["input_ids"])
        tokenized_output['token_type_ids'].append(encoded_inputs["token_type_ids"])
        tokenized_output['attention_mask'].append(encoded_inputs["attention_mask"])
        tokenized_output['mask_positions'].append(encoded_inputs["mask_position"])

        if train_mode:  # 如果是训练模式
            label_encoded = tokenizer(text=[label])  # 对标签进行编码
            label_encoded = label_encoded['input_ids'][0][1:-1]  # 去掉 [CLS] 和 [SEP]
            label_encoded = label_encoded[:max_label_len]  # 截断至最大标签长度
            label_encoded = label_encoded + [tokenizer.pad_token_id] * (max_label_len - len(label_encoded))  # 填充至最大长度
            tokenized_output['mask_labels'].append(label_encoded)  # 添加到 mask_labels 列表中

    # 根据 return_tensor 参数决定返回类型
    for k, v in tokenized_output.items():
        if return_tensor:  # 如果需要返回 Tensor 类型
            tokenized_output[k] = torch.LongTensor(v)
        else:  # 否则返回 NumPy 数组
            tokenized_output[k] = np.array(v)

    return tokenized_output  # 返回最终的 tokenized 结果


if __name__ == '__main__':
    pc = ProjectConfig()  # 创建配置对象
    train_dataset = load_dataset('text', data_files=pc.train_path)  # 加载训练数据集
    print(train_dataset)  # 打印数据集信息
    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)  # 加载预训练 Tokenizer
    hard_template = HardTemplate(prompt='这是一条{MASK}评论：{textA}')  # 创建硬模板对象

    # 定义一个偏函数，固定部分参数
    convert_func = partial(
        convert_example,
        tokenizer=tokenizer,
        hard_template=hard_template,
        max_seq_len=30,
        max_label_len=2,
    )

    # 使用 map 方法对数据集进行批量处理
    dataset = train_dataset.map(convert_func, batched=True)
    print(f"dataset-->{dataset}")  # 打印处理后的数据集

    # 打印第一个样本的信息
    for value in dataset['train']:
        print(value)  # 打印样本内容
        print(len(value['input_ids']))  # 打印 input_ids 的长度
        print(type(value['input_ids']))  # 打印 input_ids 的类型
        break  # 只打印第一个样本后退出循环
# 实现了一个数据转换流程，用于将原始文本数据转换为深度学习模型所需的输入格式。它通过 HardTemplate 类对文本进行模板化处理，并结合 Hugging Face 的 AutoTokenizer 进行编码，生成包括 input_ids、token_type_ids、attention_mask、mask_positions 和 mask_labels 在内的多种特征。该流程支持训练和推理两种模式，并可通过 map 方法高效处理大规模数据集。整体设计灵活且高效，适用于基于 Prompt 的自然语言处理任务，如文本分类、填空预测等。